package com.fy.fyai.common.config;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fy.fyai.common.ex.ConditionException;
import com.fy.fyai.common.ex.ModelNotAvailableException;
import com.fy.fyai.common.ex.handler.ExHandler;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * AI模型管理器
 *
 * @author fwq
 * @date 2025-10-25
 */
@Slf4j
@Component
public class AiModelManager {

    @Value("${spring.ai.ollama.base-url:http://localhost:11434}")
    private String ollamaBaseUrl;

    @Value("${spring.ai.ollama.chat.models}")
    private String availableModels;

    @Value("${spring.ai.ollama.chat.default-model}")
    private String defaultModelName;

    private final RestTemplate restTemplate = new RestTemplate();
    private final ObjectMapper objectMapper = new ObjectMapper();

    private List<String> availableModelList;

    @PostConstruct
    public void init() {
        // 解析可用的模型列表
        availableModelList = Arrays.stream(availableModels.split(","))
                .map(String::trim)
                .collect(Collectors.toList());
    }

    /**
     * 调用指定模型
     *
     * @param modelName 模型名称
     * @param prompt    提示词
     * @return AI响应
     */
    public String call(String modelName, String prompt) {
        validateModel(modelName);
        return callModel(modelName, prompt);
    }

    /**
     * 使用默认模型调用
     */
    public String callWithDefault(String prompt) {
        return callModel(defaultModelName, prompt);
    }

    /**
     * 调用模型（内部方法）
     */
    private String callModel(String modelName, String prompt) {
        try {
            // 使用ObjectMapper构建JSON请求体，避免手动拼接导致的特殊字符问题
            Map<String, Object> requestMap = new HashMap<>();
            requestMap.put("model", modelName);
            requestMap.put("prompt", prompt);
            requestMap.put("stream", false);
            String requestBody = objectMapper.writeValueAsString(requestMap);

            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity<String> entity = new HttpEntity<>(requestBody, headers);

            ResponseEntity<String> response = restTemplate.postForEntity(
                    ollamaBaseUrl + "/api/generate", entity, String.class
            );

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                JsonNode jsonNode = objectMapper.readTree(response.getBody());
                JsonNode responseNode = jsonNode.get("response");
                if (responseNode != null) {
                    return responseNode.asText();
                }
            }

            throw new RuntimeException("调用模型失败: " + response.getStatusCode());

        } catch (Exception e) {
            log.error("调用模型失败: {}", modelName, e);
            throw new ConditionException("调用模型失败: " + modelName);
        }
    }

    /**
     * 验证模型是否可用
     */
    private void validateModel(String modelName) {
        ExHandler.throwIf(!availableModelList.contains(modelName), String.format("模型'%s'不可用，可用模型:%s",
                modelName, String.join(", ", availableModelList)));
    }

    /**
     * 获取所有可用模型
     */
    public List<String> getAvailableModels() {
        return availableModelList;
    }

    /**
     * 获取默认模型名称
     */
    public String getDefaultModelName() {
        return defaultModelName;
    }

}